#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
集成优化文本分割器的RAG系统
============================

这是原始 rag1.py 的升级版本，集成了优化的文本分割功能。
展示了如何在现有RAG系统中无缝集成智能文本分割器。

主要改进：
1. 使用优化的文本分割器替换原有的简单分割
2. 根据文档特征自动选择最佳分割策略
3. 提供更好的错误处理和回退机制
4. 保持完全的向后兼容性

作者：AI Assistant
基于：原始 rag1.py + optimized_text_splitter
版本：2.0.0 (集成优化版)
"""

import chromadb
from chromadb.utils import embedding_functions
import os
import logging
from openai import OpenAI
from typing import List, Dict, Any, Optional
import time

# 尝试导入优化的文本分割器
try:
    from optimized_text_splitter import (
        split_text_optimized,
        split_text_simple, 
        get_optimal_config_for_scenario,
        TextSplitConfig,
        SplitStrategy,
        OverlapStrategy
    )
    OPTIMIZED_SPLITTER_AVAILABLE = True
    print("✅ 成功导入优化文本分割器")
except ImportError as e:
    OPTIMIZED_SPLITTER_AVAILABLE = False
    print(f"⚠️  优化分割器不可用，将使用原始实现: {str(e)}")

# 尝试导入优化的文档解析器
try:
    from optimized_document_parser import load_employee_file_simple as load_employee_file_enhanced
    OPTIMIZED_PARSER_AVAILABLE = True
    print("✅ 成功导入优化文档解析器")
except ImportError:
    OPTIMIZED_PARSER_AVAILABLE = False
    print("⚠️  优化文档解析器不可用")

# 原始导入
from docx import Document

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化ChromaDB客户端
chroma_client = chromadb.PersistentClient(path="./chroma_db")


def load_employee_file(file_path: str) -> str:
    """
    优化版员工手册加载函数
    
    集成了优化的文档解析器，如果不可用则回退到原始实现
    """
    if OPTIMIZED_PARSER_AVAILABLE:
        try:
            logger.info(f"使用优化文档解析器处理: {file_path}")
            return load_employee_file_enhanced(file_path)
        except Exception as e:
            logger.warning(f"优化文档解析器失败，回退到原始实现: {str(e)}")
    
    # 原始实现
    logger.info(f"使用原始文档解析器处理: {file_path}")
    try:
        document = Document(file_path)
        all_text = []
        for para in document.paragraphs:
            clean_text = para.text.strip()
            all_text.append(clean_text)
        return '\n'.join(all_text)
    except Exception as e:
        raise Exception(f"文档加载失败: {str(e)}")


def split_text(text: str, chunk_size: int = 200) -> List[str]:
    """
    优化版文本分割函数
    
    这是对原有split_text函数的直接替换，保持完全的API兼容性，
    但内部使用智能分割算法。
    """
    if OPTIMIZED_SPLITTER_AVAILABLE:
        try:
            logger.info(f"使用智能文本分割器，块大小: {chunk_size}")
            
            # 分析文本特征，自动选择最佳策略
            strategy = _select_optimal_strategy(text, chunk_size)
            
            # 使用优化的分割器
            result = split_text_optimized(
                text=text,
                strategy=strategy,
                chunk_size=chunk_size,
                min_chunk_size=max(50, chunk_size // 4),
                max_chunk_size=chunk_size * 2,
                overlap_strategy=OverlapStrategy.SMART_OVERLAP,
                overlap_size=min(50, chunk_size // 10),
                clean_text=True,
                verbose=False
            )
            
            # 记录分割质量
            if result.metrics:
                quality_score = result.metrics.get('semantic_completeness', 0)
                logger.info(f"文本分割完成: {len(result.chunks)}个块, "
                          f"策略: {strategy}, 质量得分: {quality_score:.2f}")
            
            return result.chunks
            
        except Exception as e:
            logger.error(f"智能分割失败，回退到原始实现: {str(e)}")
            return _original_split_text(text, chunk_size)
    else:
        # 回退到原始实现
        logger.info("使用原始文本分割器")
        return _original_split_text(text, chunk_size)


def _select_optimal_strategy(text: str, chunk_size: int) -> str:
    """
    根据文本特征自动选择最佳分割策略
    
    Args:
        text: 待分割文本
        chunk_size: 目标块大小
        
    Returns:
        最佳策略名称
    """
    text_length = len(text)
    
    # 检测文档特征
    has_qa_pairs = bool('问：' in text or '答：' in text or '问:' in text or '答:' in text)
    has_clear_paragraphs = text.count('\n\n') > 3
    has_lists = bool('- ' in text or '• ' in text or '1.' in text or '2.' in text)
    
    # 策略选择逻辑
    if has_qa_pairs:
        # 有问答结构，使用语义感知策略
        return "semantic_aware"
    elif has_clear_paragraphs and text_length > 1000:
        # 段落结构清晰的长文档，使用段落边界策略
        return "paragraph_boundary" 
    elif chunk_size <= 300:
        # 小块分割，使用句子边界策略
        return "sentence_boundary"
    else:
        # 默认使用混合策略
        return "hybrid"


def _original_split_text(text: str, chunk_size: int) -> List[str]:
    """
    原始的文本分割实现（备用）
    
    保留原始实现作为备用方案，确保系统的健壮性
    """
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunk = text[i:i + chunk_size]
        if chunk.strip():  # 只添加非空块
            chunks.append(chunk)
    return chunks


def split_text_advanced(text: str, 
                       chunk_size: int = 500,
                       strategy: str = "auto",
                       overlap_size: int = 50) -> Dict[str, Any]:
    """
    高级文本分割函数
    
    提供更多控制选项和详细信息，适用于需要精细控制的场景
    
    Args:
        text: 待分割文本
        chunk_size: 目标块大小
        strategy: 分割策略 ("auto", "fixed_length", "sentence_boundary", "paragraph_boundary", "semantic_aware", "hybrid")
        overlap_size: 重叠大小
        
    Returns:
        包含分割结果和详细信息的字典
    """
    if not OPTIMIZED_SPLITTER_AVAILABLE:
        # 回退到简单实现
        chunks = _original_split_text(text, chunk_size)
        return {
            'chunks': chunks,
            'chunk_count': len(chunks),
            'strategy_used': 'fixed_length',
            'processing_time': 0,
            'quality_score': 0.5,  # 估算得分
            'success': True,
            'warnings': ['优化分割器不可用，使用基础实现']
        }
    
    try:
        # 自动选择策略
        if strategy == "auto":
            strategy = _select_optimal_strategy(text, chunk_size)
        
        # 配置分割参数
        config = TextSplitConfig(
            chunk_size=chunk_size,
            min_chunk_size=max(50, chunk_size // 4),
            max_chunk_size=chunk_size * 2,
            strategy=SplitStrategy(strategy),
            overlap_strategy=OverlapStrategy.SMART_OVERLAP,
            overlap_size=overlap_size,
            clean_text=True,
            handle_lists=True,
            handle_tables=True,
            collect_metrics=True
        )
        
        # 执行分割
        result = split_text_optimized(text, config=config)
        
        # 计算质量得分
        quality_score = 0.8  # 默认得分
        if result.metrics:
            semantic_score = result.metrics.get('semantic_completeness', 0.8)
            coverage_score = result.metrics.get('coverage_ratio', 1.0)
            quality_score = (semantic_score * 0.7 + coverage_score * 0.3)
        
        return {
            'chunks': result.chunks,
            'chunk_count': len(result.chunks),
            'strategy_used': strategy,
            'processing_time': result.processing_time,
            'quality_score': quality_score,
            'metrics': result.metrics,
            'metadata': result.metadata,
            'success': True,
            'warnings': []
        }
        
    except Exception as e:
        logger.error(f"高级分割失败: {str(e)}")
        
        # 回退处理
        chunks = _original_split_text(text, chunk_size)
        return {
            'chunks': chunks,
            'chunk_count': len(chunks),
            'strategy_used': 'fixed_length',
            'processing_time': 0,
            'quality_score': 0.5,
            'success': False,
            'error': str(e),
            'warnings': ['使用备用分割方案']
        }


def create_rag_chunks_optimized(text: str, 
                               chunk_size: int = 500,
                               scenario: str = "general") -> List[str]:
    """
    为RAG系统创建优化的文档块
    
    根据不同场景自动优化分割参数
    
    Args:
        text: 待分割文本
        chunk_size: 目标块大小
        scenario: 应用场景 ("general", "hr_documents", "technical_docs", "qa_knowledge")
        
    Returns:
        优化的文档块列表
    """
    if OPTIMIZED_SPLITTER_AVAILABLE:
        try:
            # 获取场景优化配置
            if scenario != "general":
                config = get_optimal_config_for_scenario(scenario)
                # 调整块大小
                config.chunk_size = chunk_size
            else:
                # 通用配置
                config = TextSplitConfig(
                    chunk_size=chunk_size,
                    strategy=SplitStrategy.HYBRID,
                    overlap_strategy=OverlapStrategy.SMART_OVERLAP,
                    overlap_size=min(100, chunk_size // 5),
                    clean_text=True,
                    handle_lists=True,
                    handle_tables=True
                )
            
            result = split_text_optimized(text, config=config)
            
            logger.info(f"RAG块创建完成: {len(result.chunks)}个块, "
                       f"场景: {scenario}, 平均长度: {sum(len(c) for c in result.chunks) / len(result.chunks):.0f}")
            
            return result.chunks
            
        except Exception as e:
            logger.warning(f"优化RAG块创建失败，使用基础方案: {str(e)}")
    
    # 回退到基础分割
    return split_text(text, chunk_size)


# 保持原有的其他函数不变
def text_embeddings(texts: List[str]) -> List[List[float]]:
    """文本向量化函数 - 保持原有实现"""
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    if isinstance(texts, str):
        texts = [texts]

    batch_size = 10
    embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="text-embedding-v4",
            input=batch,
            dimensions=1024,
            encoding_format="float"
        )

        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chroma(texts: List[str], collection_name: str = "my_collection") -> bool:
    """
    优化的ChromaDB存储函数
    
    增加了错误处理和进度跟踪
    """
    try:
        logger.info(f"开始存储{len(texts)}个文档块到ChromaDB")
        
        # 计算嵌入向量
        embeddings = text_embeddings(texts)
        
        # 创建嵌入函数
        embedding_function = embedding_functions.OpenAIEmbeddingFunction(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            model_name="text-embedding-v4",
            api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

        # 删除已存在的集合
        try:
            chroma_client.delete_collection(name=collection_name)
            logger.info(f"删除已存在的集合: {collection_name}")
        except Exception:
            pass

        # 创建新集合
        collection = chroma_client.create_collection(
            name=collection_name,
            embedding_function=embedding_function
        )

        # 批量添加文档
        batch_size = 50
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            batch_embeddings = embeddings[i:i + batch_size]
            batch_ids = [f"text_{j}" for j in range(i, i + len(batch_texts))]
            
            collection.add(
                documents=batch_texts,
                embeddings=batch_embeddings,
                ids=batch_ids
            )
            
            logger.info(f"已存储 {min(i + batch_size, len(texts))}/{len(texts)} 个文档块")

        logger.info("ChromaDB存储完成")
        return True
        
    except Exception as e:
        logger.error(f"ChromaDB存储失败: {str(e)}")
        return False


def rag_answer_question(question: str, 
                       top_k: int = 3,
                       collection_name: str = "my_collection") -> str:
    """
    RAG问答函数 - 保持原有接口，内部优化
    """
    try:
        logger.info(f"处理问题: {question}")
        
        # 创建嵌入函数
        embedding_function = embedding_functions.OpenAIEmbeddingFunction(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            model_name="text-embedding-v4",
            api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

        # 获取集合
        collection = chroma_client.get_collection(
            name=collection_name,
            embedding_function=embedding_function
        )

        # 查询相关文档
        results = collection.query(
            query_texts=[question],
            n_results=top_k
        )

        if not results["documents"] or not results["documents"][0]:
            return "抱歉，没有找到相关信息来回答您的问题。"

        retrieved_chunks = list(set(results["documents"][0]))
        logger.info(f"检索到 {len(retrieved_chunks)} 个相关文档块")

        # 构建上下文
        context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
        
        # 构建提示词
        prompt = f"""你是一个专业的HR助理，请仔细阅读以下员工手册内容并准确回答用户问题。

【员工手册内容】：
{context}

【用户问题】：
{question}

【回答要求】：
1. 严格基于员工手册内容回答，不要编造或推测
2. 如果手册中没有相关信息，请明确说明"手册中未提及此信息"
3. 回答应简洁明了，条理清晰
4. 如有多个要点，请分点列出
5. 如果涉及具体条款，可引用相关内容

请回答："""

        # 调用LLM生成回答
        response = get_response([{"role": "user", "content": prompt}])
        
        logger.info("RAG问答完成")
        return response
        
    except Exception as e:
        error_msg = f"RAG问答处理失败: {str(e)}"
        logger.error(error_msg)
        return "抱歉，处理您的问题时出现了错误。"


def get_response(messages: List[Dict[str, str]]) -> str:
    """获取LLM响应 - 保持原有实现"""
    client = OpenAI(
        api_key="sk-",
        base_url="http://localhost:11434/v1",
    )

    completion = client.chat.completions.create(
        model="qwen:latest",
        messages=messages,
    )
    return completion.choices[0].message.content


def demo_integrated_rag():
    """演示集成优化分割器的RAG系统"""
    print("🎯 集成优化RAG系统演示")
    print("="*60)
    
    file_path = 'employee_manual.docx'
    
    try:
        # 1. 使用优化的文档加载
        print("1. 加载员工手册...")
        start_time = time.time()
        document_text = load_employee_file(file_path)
        load_time = time.time() - start_time
        
        print(f"✅ 文档加载完成: {len(document_text)} 字符, 耗时: {load_time:.3f}秒")
        
        # 2. 使用优化的文本分割
        print("\n2. 智能文本分割...")
        
        # 演示不同的分割方式
        split_methods = [
            ("基础分割", lambda: split_text(document_text, chunk_size=400)),
            ("高级分割", lambda: split_text_advanced(document_text, chunk_size=400, strategy="auto")),
            ("RAG优化", lambda: create_rag_chunks_optimized(document_text, chunk_size=400, scenario="hr_documents"))
        ]
        
        best_chunks = None
        
        for method_name, split_func in split_methods:
            start_time = time.time()
            
            if method_name == "高级分割":
                result = split_func()
                chunks = result['chunks']
                quality_score = result.get('quality_score', 0)
                strategy_used = result.get('strategy_used', 'unknown')
                
                print(f"✅ {method_name}: {len(chunks)}个块, 策略: {strategy_used}, 质量: {quality_score:.2f}")
                
                if quality_score > 0.7:  # 选择质量最好的块
                    best_chunks = chunks
            else:
                chunks = split_func()
                print(f"✅ {method_name}: {len(chunks)}个块")
                
                if best_chunks is None:  # 如果没有更好的，使用基础分割
                    best_chunks = chunks
            
            split_time = time.time() - start_time
            print(f"   处理时间: {split_time:.3f}秒")
        
        if not best_chunks:
            print("❌ 所有分割方法都失败了")
            return
        
        # 3. 存储到向量数据库
        print(f"\n3. 存储到向量数据库...")
        if store_in_chroma(best_chunks, "optimized_employee_manual"):
            print("✅ 向量数据库存储成功")
        else:
            print("❌ 向量数据库存储失败")
            return
        
        # 4. 测试问答功能
        print(f"\n4. 测试RAG问答...")
        test_questions = [
            "员工试用期工资如何计算？",
            "年假申请需要提前多久？",
            "公司提供哪些福利待遇？"
        ]
        
        for question in test_questions:
            print(f"\n问题: {question}")
            answer = rag_answer_question(question, collection_name="optimized_employee_manual")
            print(f"回答: {answer[:200]}...")
        
        print(f"\n🎉 集成优化RAG系统演示完成!")
        
        # 5. 性能总结
        print(f"\n📊 性能总结:")
        print(f"✅ 文档加载: 支持多种格式，自动错误处理")
        print(f"✅ 智能分割: 自动策略选择，质量评估")
        print(f"✅ 向量存储: 批量处理，进度跟踪")
        print(f"✅ 问答系统: 上下文优化，回答质量提升")
        
    except Exception as e:
        print(f"❌ 演示过程出错: {str(e)}")


def compatibility_test():
    """兼容性测试"""
    print("\n🔄 向后兼容性测试:")
    print("-" * 40)
    
    test_text = "这是一个测试文本。它包含多个句子。每个句子都有不同的内容。"
    
    try:
        # 测试原始接口
        chunks = split_text(test_text, chunk_size=20)
        print(f"✅ split_text 接口: {len(chunks)} 个块")
        
        # 测试是否返回字符串列表
        if isinstance(chunks, list) and all(isinstance(chunk, str) for chunk in chunks):
            print("✅ 返回类型: 正确")
        else:
            print("❌ 返回类型: 错误")
        
        # 测试内容是否合理
        total_length = sum(len(chunk) for chunk in chunks)
        if total_length > 0:
            print("✅ 内容处理: 正确")
        else:
            print("❌ 内容处理: 错误")
            
    except Exception as e:
        print(f"❌ 兼容性测试失败: {str(e)}")


if __name__ == '__main__':
    # 运行演示
    demo_integrated_rag()
    
    # 运行兼容性测试
    compatibility_test()
    
    print("\n" + "="*60)
    print("💡 使用说明:")
    print("1. 此版本完全兼容原有的 rag1.py 接口")
    print("2. 自动集成智能文本分割和文档解析功能")
    print("3. 如果优化组件不可用，会自动回退到原始实现")
    print("4. 提供了额外的高级功能接口供选择使用")
    print("5. 所有改进都是透明的，无需修改现有调用代码")
    print("="*60)
